/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { BasicBlock } from './basic_block';
import { Code } from './code';
import { StackSlot } from './stack_solt';
import { Arg } from './arg';
import { mergeIntoSet } from './util';

export class Liveness {
  private _thing: typeof StackSlot;
  get thing(): typeof StackSlot {
    return this._thing;
  }

  private _code: Code;
  get code(): Code {
    return this._code;
  }

  _liveAtHead: Map<BasicBlock, Set<StackSlot>>;
  get liveAtHead(): Map<BasicBlock, Set<StackSlot>> {
    return this._liveAtHead;
  }

  _liveAtTail: Map<BasicBlock, Set<StackSlot>>;
  get liveAtTail(): Map<BasicBlock, Set<StackSlot>> {
    return this._liveAtTail;
  }

  dirtyBlocks: Set<BasicBlock>;

  constructor(thing: typeof StackSlot, code: Code) {
    this._thing = thing;
    this._code = code;
    this._liveAtHead = new Map();
    this._liveAtTail = new Map();
    for (let block of code) {
      this._liveAtHead.set(block, new Set());

      let liveAtTail: Set<StackSlot> = new Set();
      this._liveAtTail.set(block, liveAtTail);

      block.last().forEach(thing, (value, role, type, width) => {
        if (Arg.isLateUse(role)) {
          liveAtTail.add(value);
        }
        return null;
      });
    }
    this.dirtyBlocks = new Set(code);
    this.handerCode(code, thing);
  }

  private handerCode(code: Code, thing: typeof StackSlot): void {
    let changed = false;
    do {
      changed = false;
      let blockIndex = code.size;
      while (blockIndex > 0) {
        blockIndex -= 1;
        let blockBacs = code.at(blockIndex);
        if (blockBacs === null) {
          continue;
        }
        if (!this.dirtyBlocks.delete(blockBacs)) {
          continue;
        }
        let local = this.localCalc(blockBacs);
        let instIndex = blockBacs.size;
        while (instIndex > 0) {
          instIndex -= 1;
          local.execute(instIndex);
        }
        blockBacs.at(0).forEach(thing, (value, role, type, width) => {
          if (Arg.isEarlyDef(role)) {
            local.liveSet.delete(value);
          }
          return null;
        });
        if (!mergeIntoSet(this._liveAtHead.get(blockBacs), local.liveSet)) {
          continue;
        }
        for (let predecessor of blockBacs.predecessors) {
          if (mergeIntoSet(this._liveAtTail.get(predecessor), this._liveAtHead.get(blockBacs))) {
            this.dirtyBlocks.add(predecessor);
            changed = true;
          }
        }
      }
    } while (changed);
  }

  localCalc(block: BasicBlock): LocalCalc {
    return new LocalCalc(this, block);
  }
}

export class LocalCalc {
  _liveness: Liveness;
  _block: BasicBlock;
  _liveSet: Set<StackSlot>;

  get liveSet(): Set<StackSlot> {
    return this._liveSet;
  }
  set liveSet(newValue: Set<StackSlot>) {
    this._liveSet = newValue;
  }

  constructor(liveness: Liveness, block: BasicBlock) {
    this._liveness = liveness;
    this._block = block;
    this._liveSet = liveness.liveAtTail.get(block);
  }

  execute(instIndex: number): void {
    let inst = this._block.at(instIndex);
    // First handle the early defs of the next instruction.
    if (instIndex + 1 < this._block.size) {
      this._block.at(instIndex + 1).forEach(this._liveness.thing, (value, role, type, width) => {
        if (Arg.isEarlyDef(role)) {
          this._liveSet.delete(value);
        }
        return null;
      });
    }
    // Then handle defs.
    inst.forEach(this._liveness.thing, (value, role, type, width) => {
      if (Arg.isEarlyUse(role)) {
        this._liveSet.delete(value);
      }
      return null;
    });

    // Then handle uses.
    inst.forEach(this._liveness.thing, (value, role, type, width) => {
      if (Arg.isEarlyUse(role)) {
        this._liveSet.add(value);
      }
      return null;
    });

    // Finally handle the late uses of the previous instruction.
    if (instIndex - 1 >= 0) {
      this._block.at(instIndex - 1).forEach(this._liveness.thing, (value, role, type, width) => {
        if (Arg.isLateUse(role)) {
          this._liveSet.add(value);
        }
        return null;
      });
    }
  }
}
